import numpy as np
import os
import torch
import cv2

import matplotlib.pyplot as plt


def load_data():
    pass


def imgBrightness(img1, c, b):
    rows, cols, channels = img1.shape
    blank = np.zeros([rows, cols, channels], img1.dtype)
    rst = cv2.addWeighted(img1, c, blank, 1 - c, b)
    return rst


if __name__ == '__main__':

    path = '/data/fg2021/DG-13'

    x5_path = '/data/fg2021/DG-13-x5'
    x1_path = '/data/fg2021/DG-13-x1'
    x5_path_rgb = '/data/fg2021/DG-13-x5-rgb'
    x1_path_rgb = '/data/fg2021/DG-13-x1-rgb'

    folders = os.listdir(path)

    for folder in folders:
        folders_path = path + '/' + folder
        files = os.listdir(folders_path)
        x5_path_folder = x5_path + '/' + folder
        x1_path_folder = x1_path + '/' + folder
        x5_path_rgb_folder = x5_path_rgb + '/' + folder
        x1_path_rgb_folder = x1_path_rgb + '/' + folder

        try:
            os.mkdir(x5_path_folder)
            os.mkdir(x1_path_folder)
            os.mkdir(x5_path_rgb_folder)
            os.mkdir(x1_path_rgb_folder)
        except:
            pass

        for file in files:
            name = folders_path + '/' + file
            data = np.load(name)

            for alpha in range(0, 10, 2):
                i = 0
                seq = []
                seq_rgb = []
                for frame in data:
                    i += 1
                    if i % 4 != 0:
                        continue
                    # cv2.imshow("ori",frame[:,:,0:3])
                    f = imgBrightness(frame, 1 - alpha / 10.0, 0)
                    f[:, :, 3] = frame[:, :, 3]
                    f = cv2.resize(f, (256, 256))
                    seq.append(f)
                    seq_rgb.append(f[:, :, 0:3])

                seq_np = np.array(seq).transpose(3, 0, 1, 2)
                seq_rgb_np = np.array(seq_rgb).transpose(3, 0, 1, 2)
                # print(seq_np.shape)
                # print(seq_rgb_np.shape)
                # cv2.waitKey(-1)

                np.save(x5_path_folder + '/' + file + "-" + str(alpha), seq_np)
                np.save(x5_path_rgb_folder + '/' + file + "-" + str(alpha), seq_rgb_np)

                if 0 == alpha:
                    np.save(x1_path_folder + '/' + file + "-" + str(alpha), seq_np)
                    np.save(x1_path_rgb_folder + '/' + file + "-" + str(alpha), seq_rgb_np)

        print(folder + " finish")
